import codecs
import math
import pickle
import random
import numpy as np
import torch
from torch import nn



LOAD = __import__('load')
class embedding_net(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.lut = nn.Embedding(len(item2id),LOAD.con.word_emb_size)
        self.word_embsize = LOAD.con.word_emb_size
    def forward(self,x):
        return self.lut(x)*math.sqrt(self.word_embsize)
class net_no_crf(nn.Module):
    def __init__(self):
        super().__init__()
        self.Sigmoid = nn.Sigmoid()
        self.conv1_left = nn.Conv1d(in_channels=50,out_channels=50,kernel_size=3,padding=1)
        self.conv1_right = nn.Conv1d(in_channels=50,out_channels=50,kernel_size=3,padding=1)
        self.conv2_left = nn.Conv1d(in_channels=50,out_channels=50,kernel_size=3,padding=1)
        self.conv2_right = nn.Conv1d(in_channels=50,out_channels=50,kernel_size=3,padding=1)
        self.conv3_left = nn.Conv1d(in_channels=50,out_channels=50,kernel_size=3,padding=1)
        self.conv3_right = nn.Conv1d(in_channels=50,out_channels=50,kernel_size=3,padding=1)
        self.end = nn.Conv1d(in_channels=50,out_channels=4,kernel_size=1)
    def forward(self,input):
        output1_l = self.conv1_left(input)
        output1_r = self.conv1_right(input)
        output1_r = self.Sigmoid(output1_r)
        output1 = output1_l*output1_r
        
        output2_l = self.conv2_left(output1)
        output2_r = self.conv2_right(output1)
        output1_r = self.Sigmoid(output2_r)
        output2 = output2_l*output2_r
        
        output3_l = self.conv3_left(output2)
        output3_r = self.conv3_right(output2)
        output1_r = self.Sigmoid(output3_r)
        output3 = output3_l*output3_r
        output = self.end(output3)
        return output

def test():
    net_N = torch.load('net_N.pt')
    net_emb = torch.load('net_emb.pt')
    test_data = LOAD.read_train_file(codecs.open(LOAD.con.test_path,'r','utf8'),word_window = LOAD.con.word_window)
    data_path = open('data/convertchart.data','rb')
    [item2id,id2item,tag2id,id2tag] = pickle.load(data_path)
    data_path.close()


    test_sentences = test_data[0]
    test_sen = random.choice(test_sentences)
    test_id = []
    for word in test_sen:
        if word in item2id:
            id = item2id[word]
        else:
            id = item2id['<UNK>']
        test_id.append(id)
    test_id = torch.tensor([test_id])
    test_id_emb = net_emb(test_id)
    test_id_emb = test_id_emb.permute(0,2,1)
    net_N = net_N.cpu()
    output_tags = (net_N(test_id_emb))
    predict_tags = []
    output_tags = torch.squeeze(output_tags).permute(1,0)
    for tag in output_tags:
        pre = np.argmax(tag.detach())
        pre = id2tag[pre.item()]
        predict_tags.append(pre)
    predictions = zip(test_sen,predict_tags)
    for word,tag in predictions:
        if tag == 'S' or tag == 'E':
            print(word,end=' ')
        else:
            print(word,end='')



if __name__ == '__main__':
    test()